﻿using System;
using System.Linq;
using System.Text;
using swe.Core;

namespace SqlInjectionValidator
{
	public class SqlInjectionValidator
	{
		private enum QueryState { Normal, Quote, Identifier }
		public static ValidationResult ValidateQueryFast(string query, int? stacks = 1, int? unions = null)
		{
			if (stacks.HasValue && stacks <= 0)
				throw new ArgumentException("stacks should be at least 1.");
			var currentState = QueryState.Normal;
			var i = 0;
			var len = query.Length;
			var sb = new StringBuilder(len);
			while (i >= 0 && i < len)
			{
				switch (currentState)
				{
					case QueryState.Normal:
						var newI = query.IndexOfAny("'[".ToCharArray(), i) + 1;
						if (newI!=0)
						{
							sb.Append(query.Substring(i, newI - i));
							currentState = query[newI - 1] == '\'' ? QueryState.Quote : QueryState.Identifier;
						}
						else
							sb.Append(query.Substring(i));
						i = newI;
						break;
					case QueryState.Quote:
						currentState = QueryState.Normal;
						while ( true )
						{
							i = query.IndexOf('\'', i) + 1;
							if (i >= 0 && i < len && query[i] == '\'')
								i++;
							else
								break;
						}
						break;
					case QueryState.Identifier:
						currentState = QueryState.Normal;
						i = query.IndexOf(']', i) + 1;
						break;
				}
				if (i == 0)
				{
					break;
				}
			}

			return TestForInvalidTokens(sb, stacks != null ? stacks - 1 : null, unions);
		}

		public static ValidationResult ValidateQueryUsingSwe(string query, int? stacks = 1, int? unions = null)
		{
			var len = query.Length;
			var sb = new StringBuilder(len);
			var sqlInjSm = StateMachine<int, int>.Create("Normal")
				.State("Normal", (s, i) =>
									 {
										 if (i >= len)
											 return StateMachine.Break;
										 s.UpdateObject(i + 1);
										 if (query[i] == '\'')
											 return StateMachine.NamedState("Quote");
										 if (query[i] == '[')
											 return StateMachine.NamedState("Identifier");
										 sb.Append(query[i]);
										 return StateMachine.Loop;
									 })
				.State("Quote", (s, i) =>
									{
										if (i >= len)
											return StateMachine.Break;
										if (query[i] == '\'' && i + 1 < query.Length && query[i + 1] == '\'')
										{
											s.UpdateObject(i + 2);
											return StateMachine.Loop;
										}
										s.UpdateObject(i + 1);
										return query[i] == '\'' ? StateMachine.NamedState("Normal") : StateMachine.Loop;
									})
				.State("Identifier", (s, i) =>
										 {
											 if (i >= len)
												 return StateMachine.Break;
											 s.UpdateObject(i + 1);
											 return query[i] == ']' ? StateMachine.NamedState("Normal") : StateMachine.Loop;
										 });
			var sm = sqlInjSm.BindTo(0);

			while (!sm.Eos)
			{
				sm.Process();
			}

			return TestForInvalidTokens(sb, stacks != null ? stacks - 1 : null, unions);
		}

		private static readonly string[] CommentSignature = new[] { "--", "/*" };
		private static ValidationResult TestForInvalidTokens(StringBuilder sb, int? semiCount, int? unionCount)
		{
			var q = sb.ToString();
			if (CommentSignature.Any(s => q.IndexOf(s) >= 0))
				return ValidationResult.Comment;
			if (semiCount != null)
			{
				var semicolonCnt = CountSubString(";", q);
				if (semicolonCnt > semiCount)
					return ValidationResult.Stacking;
			}
			if (unionCount != null)
			{
				var uniCount = CountSubString("UNION", q, true);
				if (uniCount > unionCount)
					return ValidationResult.UnionInjection;
			}
			return ValidationResult.NoInjectionFound;
		}

		private static int CountSubString(string sub, string str, bool ignoreCase = false)
		{
			var rv = 0;
			var pos = str.IndexOf(sub,
								  ignoreCase ? StringComparison.InvariantCultureIgnoreCase : StringComparison.Ordinal);
			while (pos >= 0)
			{
				rv++;
				pos = str.IndexOf(sub, pos + 1,
								  ignoreCase ? StringComparison.InvariantCultureIgnoreCase : StringComparison.Ordinal);
			}
			return rv;
		}
	}
}
